import json
import sqlite3
import numpy as np

class CoCoMetricDB:
    def __init__(self, db_path):
        try:
            self.conn = sqlite3.connect(db_path)
            self.cursor = self.conn.cursor()
            self._create_tables()
        except sqlite3.Error as e:
            print(f"An error occurred when connecting to the database: {e}")

    def _create_tables(self):
        try:
            self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS images (
                id INTEGER PRIMARY KEY,
                image_id INTEGER,
                ori_shape TEXT,
                scale_factor TEXT,
                img_path TEXT,
                pad_shape TEXT,
                img_shape TEXT,
                batch_input_shape TEXT,
                dt_bboxs TEXT,
                dt_segs TEXT,
                dt_labels TEXT,
                dt_scores TEXT,
                gt_bboxs TEXT,
                gt_segs TEXT,
                gt_labels TEXT
            )
            ''')

            self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS eval_match_pairs (
                id INTEGER PRIMARY KEY,
                image_id INTEGER,
                category_id INTEGER,
                area TEXT,
                iou_threshold REAL,
                annotation_type TEXT,
                dt_id INTEGER,
                gt_id INTEGER,
                dt_bbox TEXT,
                gt_bbox TEXT,
                dt_segs TEXT,
                gt_segs TEXT,
                score REAL
            )
            ''')

            self.conn.commit()
        except sqlite3.Error as e:
            print(f"An error occurred when creating tables: {e}")

    def save_annotations(self, data_samples):
        try:
            for sample in data_samples:
                # print('sample = ', sample)
                # print('sample["pred_instances"]["bboxes"] = ', sample["pred_instances"]["bboxes"])
                self.cursor.execute('''
                INSERT INTO images (
                    image_id, ori_shape, scale_factor, img_path, 
                    pad_shape, img_shape, batch_input_shape, 
                    dt_bboxs, dt_labels, dt_scores, 
                    gt_bboxs, gt_labels
                ) 
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    sample["img_id"], 
                    str(sample["ori_shape"]), str(sample["scale_factor"]), sample["img_path"],
                    str(sample["pad_shape"]), str(sample["img_shape"]), str(sample["batch_input_shape"]),
                    str(sample["pred_instances"]["bboxes"]), str(sample["pred_instances"]["labels"]), str(sample["pred_instances"]["scores"]),
                    str(sample["gt_instances"]["bboxes"]), str(sample["gt_instances"]["labels"])
                ))

            self.conn.commit()
        except sqlite3.Error as e:
            print(f"An error occurred when saving annotations: {e}")

    def save_match_pairs(self, per_image_results):
        try:
            for img_id, categories in per_image_results.items():
                for category_id, areas in categories.items():
                    for area, iou_data in areas.items():
                        for iou_threshold, results in iou_data.items():
                            self._insert_match_pairs(img_id, category_id, area, iou_threshold, 'TP', results["TP"])
                            self._insert_match_pairs(img_id, category_id, area, iou_threshold, 'FP', results["FP"])
                            self._insert_match_pairs(img_id, category_id, area, iou_threshold, 'FN', results["FN"])
            self.conn.commit()
        except sqlite3.Error as e:
            print(f"An error occurred when saving data: {e}")

    def _get_area_category(self, area):
        if area == (0, 1024):
            return 'small'
        elif area == (1024, 9216):
            return 'medium'
        elif area == (9216, 10000000000.0):
            return 'large'
        elif area == (0, 10000000000.0):
            return 'all'
        else:
            return 'unknown'
        
    def _insert_match_pairs(self, image_db_id, category_id, area, iou_threshold, annotation_type, annotations):
        try:
            for anno in annotations:
                # print('anno = ', anno)
                area_code = self._get_area_category(area)

                # Convert numpy array to list and then to JSON string
                if isinstance(anno['dt_bbox'], np.ndarray):
                    dt_bbox_list = anno['dt_bbox'].tolist()
                    dt_bbox_str = json.dumps(dt_bbox_list)
                else:
                    dt_bbox_str = json.dumps(anno['dt_bbox'])
                
                self.cursor.execute('''
                INSERT INTO eval_match_pairs (
                    image_id, category_id, area, iou_threshold, 
                    annotation_type, dt_id, gt_id, dt_bbox, gt_bbox, score
                ) 
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    image_db_id, category_id, area_code, iou_threshold, annotation_type, 
                    anno['dt_id'], anno['gt_id'], dt_bbox_str, str(anno['gt_bbox']), anno['score']
                ))
        except sqlite3.Error as e:
            print(f"An error occurred when inserting annotations: {e}")

    def get_eval_details(self, filters):
        query = "SELECT * FROM eval_match_pairs"
        conditions = []
        params = []

        if filters:
            if 'category_id' in filters:
                conditions.append("category_id IN (%s)" % ','.join('?' * len(filters['category_id'])))
                params.extend(filters['category_id'])

            if 'area' in filters:
                conditions.append("area IN (%s)" % ','.join('?' * len(filters['area'])))
                params.extend(filters['area'])

            if 'iou_threshold' in filters:
                conditions.append("iou_threshold IN (%s)" % ','.join('?' * len(filters['iou_threshold'])))
                params.extend(filters['iou_threshold'])

            if conditions:
                query += " WHERE " + " AND ".join(conditions)

        try:
            self.cursor.execute(query, params)
            return self.cursor.fetchall()
        except sqlite3.Error as e:
            print(f"An error occurred during query execution: {e}")
            return []
    
    def close(self):
        try:
            if self.conn:
                self.conn.close()
        except sqlite3.Error as e:
            print(f"An error occurred when closing the database: {e}")
